
import sys
import cv2
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QLabel, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, QLineEdit, QGridLayout, QGroupBox, QMessageBox, QComboBox
)
from PyQt6.QtCore import Qt, QTimer
from PyQt6.QtGui import QPixmap, QImage, QColor, QFont
from yolov5.infer import Infer


class CameraWindow(QWidget):
    """摄像头画面显示窗口"""
    def __init__(self, input1, input2, job_type1, job_type2, pt_file1, pt_file2, is_image=False):
        super().__init__()

        self.setWindowTitle("安全巡检")
        self.resize(1200, 800)

        # 保存输入、作业类型和模型文件路径
        self.input1 = input1
        self.input2 = input2
        self.job_type1 = job_type1
        self.job_type2 = job_type2
        self.pt_file1 = pt_file1
        self.pt_file2 = pt_file2
        self.is_image = is_image

        if is_image:
            self.cap1 = None
            self.cap2 = None
        else:
            # 初始化摄像头流
            self.cap1 = cv2.VideoCapture(input1)
            self.cap2 = cv2.VideoCapture(input2)

        # 主布局
        main_layout = QHBoxLayout()

        # 推理模型实例
        self.model_left = Infer(pt_file1, job_type1)
        self.model_right = Infer(pt_file2, job_type2)

        # 左侧摄像头 1 布局
        camera1_layout = QVBoxLayout()
        self.camera1_label = QLabel("摄像头 1")
        self.camera1_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.camera1_label.setStyleSheet("background-color: black; color: white; font-size: 16px; border-radius: 10px;")
        self.camera1_label.setFixedHeight(500)
        self.camera1_info = QLabel(f"模型文件: {pt_file1}")
        self.camera1_info.setStyleSheet("background-color: #f0f0f0; padding: 5px; border-radius: 5px;")
        
        # 检测结果显示
        self.camera1_detection_info = QLabel("检测结果：无目标")
        self.camera1_detection_info.setStyleSheet("background-color: #f0f0f0; padding: 5px; border-radius: 5px;")
        
        # 告警信息
        self.camera1_alert = QLabel("无警报")
        self.camera1_alert.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.camera1_alert.setStyleSheet("background-color: #4CAF50; color: white; font-size: 14px; font-weight: bold; padding: 5px; border-radius: 5px;")
        
        # 按顺序添加：检测结果 -> 告警信息
        camera1_layout.addWidget(self.camera1_label)
        camera1_layout.addWidget(self.camera1_info)
        camera1_layout.addWidget(self.camera1_detection_info)
        camera1_layout.addWidget(self.camera1_alert)

        # 右侧摄像头 2 布局
        camera2_layout = QVBoxLayout()
        self.camera2_label = QLabel("摄像头 2")
        self.camera2_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.camera2_label.setStyleSheet("background-color: black; color: white; font-size: 16px; border-radius: 10px;")
        self.camera2_label.setFixedHeight(500)
        self.camera2_info = QLabel(f"模型文件: {pt_file2}")
        self.camera2_info.setStyleSheet("background-color: #f0f0f0; padding: 5px; border-radius: 5px;")
        
        # 检测结果显示
        self.camera2_detection_info = QLabel("检测结果：无目标")
        self.camera2_detection_info.setStyleSheet("background-color: #f0f0f0; padding: 5px; border-radius: 5px;")
        
        # 告警信息
        self.camera2_alert = QLabel("无警报")
        self.camera2_alert.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.camera2_alert.setStyleSheet("background-color: #4CAF50; color: white; font-size: 14px; font-weight: bold; padding: 5px; border-radius: 5px;")
        
        # 按顺序添加：检测结果 -> 告警信息
        camera2_layout.addWidget(self.camera2_label)
        camera2_layout.addWidget(self.camera2_info)
        camera2_layout.addWidget(self.camera2_detection_info)
        camera2_layout.addWidget(self.camera2_alert)

        # 将左右布局加入主布局
        main_layout.addLayout(camera1_layout)
        main_layout.addLayout(camera2_layout)
        self.setLayout(main_layout)

        # 定时器用于更新摄像头画面
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_frames)
        self.timer.start(30)


    def update_frames(self):
        """更新摄像头画面"""
        if self.is_image:
            frame1 = cv2.imread(self.input1)
            frame2 = cv2.imread(self.input2)

            if frame1 is not None:
                self.model_left.infer_image(frame1)
                detection_info1 = self.model_left.parse_result()
                alert1 = self.model_left.get_alert_info()
                annotated_frame = self.model_left.plot_pred()
                self.display_frame(annotated_frame, self.camera1_label)
                self.update_detection_info(self.camera1_detection_info, detection_info1)
                self.update_alert_info(self.camera1_alert, alert1)

            if frame2 is not None:
                self.model_right.infer_image(frame2)
                detection_info2 = self.model_right.parse_result()
                alert2 = self.model_right.get_alert_info()
                annotated_frame = self.model_right.plot_pred()
                self.display_frame(annotated_frame, self.camera2_label)
                self.update_detection_info(self.camera2_detection_info, detection_info2)
                self.update_alert_info(self.camera2_alert, alert2)
        else:
            if self.cap1 and self.cap1.isOpened():
                ret1, frame1 = self.cap1.read()
                if ret1:
                    self.model_left.infer_image(frame1)
                    detection_info1 = self.model_left.parse_result()
                    alert1 = self.model_left.get_alert_info()
                    annotated_frame = self.model_left.plot_pred()
                    self.display_frame(annotated_frame, self.camera1_label)
                    self.update_detection_info(self.camera1_detection_info, detection_info1)
                    self.update_alert_info(self.camera1_alert, alert1)

            if self.cap2 and self.cap2.isOpened():
                ret2, frame2 = self.cap2.read()
                if ret2:
                    self.model_right.infer_image(frame2)
                    detection_info2 = self.model_right.parse_result()
                    alert2 = self.model_right.get_alert_info()
                    annotated_frame = self.model_right.plot_pred()
                    self.display_frame(annotated_frame, self.camera2_label)
                    self.update_detection_info(self.camera2_detection_info, detection_info2)
                    self.update_alert_info(self.camera2_alert, alert2)

    def update_alert_info(self, alert_label, alert_info):
        """根据告警信息更新告警标签"""
        if alert_info:
            if alert_info == 1:
                alert_label.setText("警报：有人员误入")
            elif alert_info == 2:
                alert_label.setText("警报：未安全着装")
            elif alert_info == 3:
                alert_label.setText("警报：没有灭火器")
            elif alert_info == 4:
                alert_label.setText("警报：没有水桶")
            elif alert_info == 5:
                alert_label.setText("警报：灭火器距离动火点太远")
            elif alert_info == 6:
                alert_label.setText("警报：水桶距离动火点太远")
            alert_label.setStyleSheet("background-color: #F44336; color: white; font-size: 14px; font-weight: bold; padding: 5px; border-radius: 5px;")  # 红色背景表示警报
        else:
            alert_label.setText("无警报")
            alert_label.setStyleSheet("background-color: #4CAF50; color: white; font-size: 14px; font-weight: bold; padding: 5px; border-radius: 5px;")  # 绿色背景表示正常

    def display_frame(self, frame, label):
        """将帧显示在指定的 QLabel 上"""
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        height, width, channel = frame.shape
        qimg = QImage(frame.data, width, height, channel * width, QImage.Format.Format_RGB888)
        pixmap = QPixmap.fromImage(qimg)
        label.setPixmap(pixmap.scaled(label.width(), label.height(), Qt.AspectRatioMode.KeepAspectRatio))
    
    def update_detection_info(self, label, detection_info):
        """更新检测信息标签"""
        if detection_info:
            label.setText("检测结果：\n" + ", ".join(detection_info))
        else:
            label.setText("检测结果：无目标")

    def closeEvent(self, event):
        """窗口关闭事件，释放资源"""
        self.release_resources()
        event.accept()

    def release_resources(self):
        """释放资源"""
        if self.cap1:
            self.cap1.release()
        if self.cap2:
            self.cap2.release()


class MainWindow(QMainWindow):
    """主界面"""
    def __init__(self):
        super().__init__()

        self.setWindowTitle("智慧园区安全预警APP")
        self.resize(900, 600)

        # 保存模型文件路径
        self.pt_file1 = ""
        self.pt_file2 = ""

        # 主布局
        main_layout = QVBoxLayout()

        # Logo 标题部分
        logo_group = QGroupBox()
        logo_group.setStyleSheet("background-color:rgb(220, 220, 220); border-radius: 10px; padding: 10px;")
        logo_layout = QHBoxLayout()
        logo_group.setLayout(logo_layout)

        # Logo
        self.logo_label = QLabel()
        self.pixmap = QPixmap('logo.png')
        self.pixmap = self.pixmap.scaled(150, 150, Qt.AspectRatioMode.KeepAspectRatio)
        self.logo_label.setPixmap(self.pixmap)
        logo_layout.addWidget(self.logo_label, alignment=Qt.AlignmentFlag.AlignLeft)

        # 标题
        self.title_label = QLabel("智慧园区安全预警APP")
        self.title_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.title_label.setFont(QFont("Arial", 24, QFont.Weight.Bold))
        logo_layout.addWidget(self.title_label, alignment=Qt.AlignmentFlag.AlignRight)

        # RTSP 地址部分
        rtsp_group = QGroupBox()
        rtsp_group.setStyleSheet("background-color:rgb(100, 149, 237); border-radius: 10px; padding: 10px;")
        rtsp_layout = QGridLayout()
        rtsp_group.setLayout(rtsp_layout)
        self.subtitle_label0 = QLabel("任务视角选择")
        self.subtitle_label0.setStyleSheet("font-size: 16px; font-weight: bold; color: white;")
        self.rtsp1_input = QLineEdit()
        self.rtsp1_input.setPlaceholderText("输入摄像头 1 的 RTSP 地址")
        self.rtsp2_input = QLineEdit()
        self.rtsp2_input.setPlaceholderText("输入摄像头 2 的 RTSP 地址")
        self.rtsp1_input.setStyleSheet("padding: 10px; font-size: 14px; border-radius: 5px; border: 2px solid #ccc; background-color: white;")
        self.rtsp2_input.setStyleSheet("padding: 10px; font-size: 14px; border-radius: 5px; border: 2px solid #ccc; background-color: white;")
        # self.rtsp1_input.setText("rtsp://admin:bjkjdx12@192.168.31.101:554")
        self.rtsp1_input.setText("D:/PythonProjects/Datasets/H-test/b (3).jpg")
        # self.rtsp2_input.setText("rtsp://admin:bjkjdx12@192.168.31.102:554")
        self.rtsp2_input.setText("D:/PythonProjects/Datasets/H-test/c (4).jpg")
        rtsp_layout.addWidget(self.subtitle_label0, 0, 0, 1, 4, alignment=Qt.AlignmentFlag.AlignCenter)
        rtsp_layout.addWidget(QLabel("摄像头 1:"), 1, 0)
        rtsp_layout.addWidget(self.rtsp1_input, 1, 1)
        rtsp_layout.addWidget(QLabel("摄像头 2:"), 1, 2)
        rtsp_layout.addWidget(self.rtsp2_input, 1, 3)

        # 添加摄像头1的作业类型选择下拉框
        self.rtsp1_job_type = QComboBox()
        self.rtsp1_job_type.addItems(["吊装作业", "动火作业"])
        self.rtsp1_job_type.setStyleSheet("""
            QComboBox {
                background-color: white;
                color: black;
                border: 2px solid #ccc;
                border-radius: 5px;
            }
            QComboBox QAbstractItemView {
                background-color: white;
                color: black;
                selection-background-color: gray;
                selection-color: black;
            }
        """)
        rtsp_layout.addWidget(QLabel("作业类型:"), 2, 0)
        rtsp_layout.addWidget(self.rtsp1_job_type, 2, 1)

        # 添加摄像头2的作业类型选择下拉框
        self.rtsp2_job_type = QComboBox()
        self.rtsp2_job_type.addItems(["吊装作业", "动火作业"])
        self.rtsp2_job_type.setStyleSheet("""
            QComboBox {
                background-color: white;
                color: black;
                border: 2px solid #ccc;
                border-radius: 5px;
            }
            QComboBox QAbstractItemView {
                background-color: white;
                color: black;
                selection-background-color: gray;
                selection-color: black;
            }
        """)
        rtsp_layout.addWidget(QLabel("作业类型:"), 2, 2)
        rtsp_layout.addWidget(self.rtsp2_job_type, 2, 3)


        # 模型文件选择部分
        pt_group = QGroupBox()
        pt_group.setStyleSheet("background-color:rgb(255, 192, 203); border-radius: 10px; padding: 10px;")
        pt_layout = QGridLayout()
        pt_group.setLayout(pt_layout)
        self.subtitle_label1 = QLabel("模型选择")
        self.subtitle_label1.setStyleSheet("font-size: 16px; font-weight: bold; color: white;")
        pt_layout.addWidget(self.subtitle_label1, 0, 0, 1, 2, alignment=Qt.AlignmentFlag.AlignCenter)
        self.pt1_button = QPushButton("选择模型 1")
        self.pt1_button.clicked.connect(self.select_pt1)
        self.pt1_button.setStyleSheet("background-color: #4CAF50; color: white; padding: 10px; border-radius: 5px; font-size: 14px;")
        self.pt1_label = QLabel("未选择模型")
        self.pt2_button = QPushButton("选择模型 2")
        self.pt2_button.clicked.connect(self.select_pt2)
        self.pt2_button.setStyleSheet("background-color: #4CAF50; color: white; padding: 10px; border-radius: 5px; font-size: 14px;")
        self.pt2_label = QLabel("未选择模型")
        pt_layout.addWidget(self.pt1_button, 1, 0)
        pt_layout.addWidget(self.pt1_label, 2, 0)
        pt_layout.addWidget(self.pt2_button, 1, 1)
        pt_layout.addWidget(self.pt2_label, 2, 1)

        # 启动按钮
        self.start_button = QPushButton("启动摄像头画面")
        self.start_button.setStyleSheet("background-color: #008CBA; color: white; padding: 15px; border-radius: 10px; font-size: 16px;")
        self.start_button.clicked.connect(self.start_action)

        # 主界面布局
        main_layout.addWidget(logo_group)
        main_layout.addWidget(rtsp_group)
        main_layout.addWidget(pt_group)
        main_layout.addWidget(self.start_button, alignment=Qt.AlignmentFlag.AlignCenter)

        # 设置主控界面
        central_widget = QWidget()
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)

    def select_pt1(self):
        """选择第一个模型文件"""
        file_path, _ = QFileDialog.getOpenFileName(self, "选择模型文件 1", "model", "PyTorch 模型文件 (*.pt)")
        if file_path:
            self.pt_file1 = file_path
            self.pt1_label.setText(file_path)

    def select_pt2(self):
        """选择第二个模型文件"""
        file_path, _ = QFileDialog.getOpenFileName(self, "选择模型文件 2", "model", "PyTorch 模型文件 (*.pt)")
        if file_path:
            self.pt_file2 = file_path
            self.pt2_label.setText(file_path)

    def start_action(self):
        """启动按钮的逻辑处理"""
        input1 = self.rtsp1_input.text()
        input2 = self.rtsp2_input.text()

        # 获取摄像头1的作业类型
        job_type1 = self.rtsp1_job_type.currentText()
        # 获取摄像头2的作业类型
        job_type2 = self.rtsp2_job_type.currentText()

        # 定义作业类型到整数值的映射
        job_type_mapping = {
            "吊装作业": 0,
            "动火作业": 1
        }

        # 根据选择的作业类型获取整数值
        job_type1_int = job_type_mapping.get(job_type1, -1)
        job_type2_int = job_type_mapping.get(job_type2, -1)

        if job_type1_int == -1 or job_type2_int == -1:
            self.show_error_message("错误", "请选择有效的作业类型")
            return

        # 检查输入和模型文件是否输入
        if not input1 or not input2:
            self.show_error_message("输入不能为空", "请填写两个输入")
            return
        if not self.pt_file1 or not self.pt_file2:
            self.show_error_message("模型文件不能为空", "请为两个摄像头选择模型文件")
            return

        # 判断输入是图片还是 RTSP 流
        is_image = input1.lower().endswith(('.png', '.jpg', '.jpeg')) and input2.lower().endswith(('.png', '.jpg', '.jpeg'))

        # 打开摄像头显示窗口，并传递模型文件路径
        self.camera_window = CameraWindow(input1, input2, job_type1_int, job_type2_int, self.pt_file1, self.pt_file2, is_image)
        self.camera_window.show()

    def show_error_message(self, title, message):
        """显示错误信息弹窗"""
        msg_box = QMessageBox(self)
        msg_box.setIcon(QMessageBox.Icon.Critical)
        msg_box.setWindowTitle(title)
        msg_box.setText(message)
        msg_box.exec()


# 主程序运行
if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())
